import numpy as np  
from PIL import Image  
from scipy.ndimage import gaussian_filter  
  
def ssim(img1, img2):  
    """  
    计算两幅灰度图像的SSIM（结构相似性度量）。  
  
    参数:  
    img1, img2 (numpy arrays): 输入的两幅灰度图像，数据类型应为float且值域在[0, 1]之间。  
  
    返回:  
    float: SSIM值，范围在[-1, 1]之间，值越大表示图像越相似。  
    """  
    C1 = (0.01 * 255)**2  
    C2 = (0.03 * 255)**2  
    eps = 1e-10  # 添加一个小的正数以防止除以零  
  
    img1 = img1.astype(np.float64)  
    img2 = img2.astype(np.float64)  
    K1 = 0.01  
    K2 = 0.03  
    L = 255  # 灰度级别，对于8位灰度图像  
  
    # 计算均值  
    mu1 = gaussian_filter(img1, 11.0)  
    mu2 = gaussian_filter(img2, 11.0)  
  
    # 计算方差和协方差  
    mu1_sq = mu1 * mu1  
    mu2_sq = mu2 * mu2  
    mu1_mu2 = mu1 * mu2  
    sigma1_sq = gaussian_filter(img1 * img1, 11.0) - mu1_sq  
    sigma2_sq = gaussian_filter(img2 * img2, 11.0) - mu2_sq  
    sigma12 = gaussian_filter(img1 * img2, 11.0) - mu1_mu2  
  
    # 计算SSIM  
    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2 + eps))  
  
    # 返回整个图像的平均值作为SSIM值  
    return np.mean(ssim_map)  
  
def preprocess_image(img_path):  
    """  
    加载图像，转换为灰度图，并归一化到[0, 1]范围。  
    """  
    img = Image.open(img_path).convert('L')  # 转换为灰度图  
    img = np.array(img, dtype=np.float64) / 255.0  # 归一化到[0, 1]范围  
    return img  
  
# 示例：使用SSIM函数  
img1_path = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\5.png'  # 替换为你的第一张图像路径  
img2_path = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\ans.png'  # 替换为你的第二张图像路径  
  
img1 = preprocess_image(img1_path)  
img2 = preprocess_image(img2_path)  
  
# 确保两张图像尺寸相同  
assert img1.shape == img2.shape, "Images must have the same dimensions."  
  
# 计算SSIM值  
ssim_value = ssim(img1, img2)  
  
# 打印SSIM值  
print(f"The SSIM value between the two images is: {ssim_value}")